import os
import numpy as np
import joblib
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.impute import SimpleImputer
from sklearn.metrics import mean_squared_error, roc_curve, auc, classification_report
from sklearn.preprocessing import StandardScaler, LabelBinarizer, LabelEncoder
"""
测试文件
将运行MODEL_NAME处训练好的模型装载，预测dic的数据
"""

HOUSING_PATH = "train/datasets/kddcup"
MODEL_NAME = "train/test3/rdm_forest_clf.pkl"
dic = {'duration': [0], 'protocol_type': 'tcp', 'service': 'http', 'flag': 'SF', 'src_bytes': 181,
       'dst_bytes': 5450, 'land': 0, 'wrong_fragment': 0, 'urgent': 0, 'ho': 0, 'num_failed_logins': 0, 'logged_in': 1,
       'num_compromised': 0, 'root_shell': 0, 'su_attempted': 0, 'num_root': 0, 'num_file_creations': 0,
       'num_shells': 0, 'num_access_files': 0, 'num_outbound_cmds': 0, 'is_host_login': 0, 'is_guest_login': 0,
       'count': 8, 'srv_count': 8, 'serror_rate': 0.0, 'srv_serror_rate': 0.0, 'rerror_rate': 0.0,
       'srv_rerror_rate': 0.0, 'same_srv_rate': 1.0, 'diff_srv_rate': 0.0, 'srv_diff_host_rate': 0.0,
       'dst_host_count': 9, 'dst_host_srv_count': 9, 'dst_host_same_srv_rate': 1.0, 'dst_host_diff_srv_rate': 0.0,
       'dst_host_same_src_port_rate': 0.11, 'dst_host_srv_diff_host_rate': 0.0, 'dst_host_serror_rate': 0.0,
       'dst_host_srv_serror_rate': 0.0, 'dst_host_rerror_rate': 0.0, 'dst_host_srv_rerror_rate': 0.0,
       'class': 'normal.'}


def load_dataset(housing_path=HOUSING_PATH):
    csv_path = os.path.join(housing_path, "kddcup.data_10_percent_corrected")
    return pd.read_csv(csv_path)

class DataFrameSelector(BaseEstimator, TransformerMixin):
    def __init__(self, attribute_names):
        self.attribute_names = attribute_names

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        return X[self.attribute_names].values

    def fit_transform(self, X, y=None):
        return self.transform(X)

# data的cat_attrib属性变为独热码返回，行数不变，列数变为此属性的种类数
def getOneHotCode(data, cat_attrib, whole_dataset):
    selector = DataFrameSelector(cat_attrib)
    cat_1 = selector.transform(data)
    encoder = LabelBinarizer()
    tmp = selector.transform(whole_dataset)
    encoder.fit(tmp)
    tmp = encoder.transform(cat_1)
    # print(encoder.classes_)
    return tmp



def getLabelEncode(data, cat_attrib, whole_dataset):
    selector = DataFrameSelector(cat_attrib)
    cat_1 = selector.transform(data)
    encoder = LabelEncoder()
    tmp = selector.transform(whole_dataset)
    encoder.fit(tmp.ravel())
    tmp = encoder.transform(cat_1.ravel())
    # print(encoder.classes_)
    for ele in range(len(encoder.classes_)):
        # print(encoder.classes_[ele])
        if encoder.classes_[ele] == "normal.":
            q_NORMAL_NUM = ele
            break
    # print(type(encoder.classes_))
    return tmp, q_NORMAL_NUM


def full_pipeline(data, num_attribs, cat_attribs, whole_dataset):
    selector = DataFrameSelector(num_attribs)
    num_1 = selector.transform(data)
    imputer = SimpleImputer(strategy="median")
    num_2 = imputer.fit_transform(num_1)
    std_scaler = StandardScaler()
    num_3 = std_scaler.fit_transform(num_2)
    # print(num_3.shape)
    # print(num_3)

    for ele in cat_attribs:
        tmp = getOneHotCode(data, ele, whole_dataset)
        # print(type(tmp))
        # print(tmp.shape)
        # df_ = pd.DataFrame(tmp, columns=list(label_binarizer.classes_))
        # print(df_.head())
        num_3 = np.c_[num_3, tmp]
    return num_3

if __name__ == '__main__':

    # X = [0,'tcp','http','SF',181,5450,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,8,8,0.0,0.0, 0.0,0.0,1.0,0.0,0.0,9,9,1.0,0.0,0.11,0.0,0.0,0.0,0.0,0.0,'normal.']

    X = pd.DataFrame(dic)  # (1, 42)

    NORMAL_NUM = -1
    kddcup = load_dataset()  # <class 'pandas.core.frame.DataFrame'>

    labels, NORMAL_NUM = getLabelEncode(kddcup, ["class"], kddcup)
    print("NORMAL_NUM = ", NORMAL_NUM)
    y_labels, _ = getLabelEncode(X, ["class"], kddcup)
    y_labels = np.array(y_labels, dtype=float)
    print("y_labels = ", y_labels)
    kddcup["class"] = labels

    cat_attribs = ["protocol_type", "service", "flag"]  # 3
    num_attribs = list(kddcup)  # 42 - 3 - 1 = 38
    num_attribs.remove("class")
    for ele in cat_attribs:
        num_attribs.remove(ele)
    # print(kddcup["class"].value_counts())

    set_labels = (labels == NORMAL_NUM)
    set = kddcup.drop("class", axis=1)  # (494021, 41)
    X_labels = (y_labels == NORMAL_NUM)

    X_set = X.drop("class", axis=1)  # (1, 41)
    print(X_set.shape)

    # (395216, 118) (98805, 118)
    set_prepared = full_pipeline(kddcup, num_attribs, cat_attribs, kddcup)  # (494021, 118)
    print(set_prepared.shape)
    test_set_prepared = full_pipeline(X_set, num_attribs, cat_attribs, kddcup)  # (1, 118)
    print(test_set_prepared.shape)


    X = []
    # 导入模型
    clf = joblib.load(MODEL_NAME)
    result = clf.predict(test_set_prepared)
    print("result:", result)
    print("labels:", X_labels)
